#!/usr/bin/env ipython

import os
import sys
import random

sys.path.append('../../lib')
sys.path.append(os.getenv("HOME")+'/lib')

from flgrcore import *
from flgrarith import *
from flgrmorpho import *
from flgrdatatools import *

### PROCEDURES #######################################################
def buildColorTable():
    r=[]
    g=[]
    b=[]
    for i in range(256):
        r.append(255-i)
        g.append(255-i)
        b.append(255-i)
    random.shuffle(r)
    random.shuffle(g)
    random.shuffle(b)
    return r,g,b

def constructPriorityStack(imlabel, imin, minGray=0, maxGray=255):
    pstack = flgr_pstack_create(imlabel.type)

    for i in range(imlabel.size_y):
        for j in range(imlabel.size_x):
            if (imlabel[i][j]>=minGray) and (imlabel[i][j]<=maxGray):
                flgr_pstack_push(pstack, imin[i][j], i, j)

    return pstack

def fillImageWithPStack(imout,pstack):
    stack = flgr_pstack_get_first_no_empty_stack(pstack)
    for i in range(pstack.nbPriority):

        color = flgr_pstack_get_label_stack(stack)
        while(stack.size>0):
            pixel=flgr_pstack_pop_stack(stack)
            row = pixel[1]
            col = pixel[2]
            imout[row][col]=color
        stack = stack.next
        

def watershed(imlabel,imin,nhb):
    cont = TNeighborBox(nhb)
    pstack = constructPriorityStack(imlabel,imin,1,255)
    
    stack = flgr_pstack_get_first_no_empty_stack(pstack)
    i=0
    while i<pstack.nbPriority:

        color = flgr_pstack_get_label_stack(stack)
        #print 'Stack priority=',color,' / stack size=',stack.size,' / Number of stack',pstack.nbPriority
        
        while(stack.size>0):
            pixel=flgr_pstack_pop_stack(stack)
            row = pixel[1]
            col = pixel[2]
            flgr2d_get_neighborhood(cont,imlabel,nhb,row,col)
 
            for k in range(cont.size):
                m = cont[0][k] + row - cont.size_y/2
                n = cont[1][k] + col - cont.size_x/2
                if cont[2][k]==0:
                    imlabel[m][n]=imlabel[row][col]
                    newcolor = imin[m][n]
                    if newcolor < color:
                        if flgr_pstack_is_last_stack(stack):
                            newcolor = color
                        else:
                            newcolor = flgr_pstack_get_label_stack(stack.next)
                    flgr_pstack_push(pstack, newcolor, m, n)
        stack = stack.next
        i=i+1



### MAIN #############################################################

imin           = flgr2d_load_pgm(os.getenv("HOME")+'/images/Gray/wmroad2.pgm','input')
imfiltered     = FLGR_Data2D(imin,"filter")
immarker       = flgr2d_load_pgm(os.getenv("HOME")+'/images/Gray/wmroad2_marker.pgm','input')
imlabel        = FLGR_Data2D(imin,"out")
imlabel_native = FLGR_Data2D(imin,"out_native")
nhb            = FLGR_Data2D(3,3,imin.type,'sq3x3','RECT','8-connexity')

flgr2d_gradient_thick(imfiltered,imin,nhb)

flgr2d_label(imlabel,immarker,nhb.connexity)
flgr2d_label(imlabel_native,immarker,nhb.connexity)

print "python watershed"
watershed(imlabel,imin,nhb)
print "C native watershed"
flgr2d_watershed(imlabel_native,imin,nhb.connexity)

flgr2d_save_pgm(imlabel)
flgr2d_save_pgm(imlabel_native)
